import pickle
from scripts.server.chat_api import chat
from tqdm import tqdm

with open('../entities.pkl', 'rb') as f:
    entities = pickle.load(f)

SUMMARIZE_PROMPT = """
您是一位乐于助人的助手，负责生成下面提供的数据的综合摘要。
给定一个或两个实体以及一系列描述，它们都与同一实体或实体组相关。
请将所有这些连接成一个单一的、全面的描述。确保包含从所有描述中收集的信息。
如果提供的描述相互矛盾，请解决矛盾并提供一个单一的、连贯的摘要。
确保以第三人称书写，并包含实体名称，以便我们了解完整的上下文。

#######
-Data-
实体：{entity_name}
描述列表：{description_list}
#######
输出：
"""


if __name__ == '__main__':
    # 合并entities中同名称的描述
    entities_new = {}
    for i, entity in enumerate(entities):
        entity_name, entity_type, entity_desc = entity
        entity_desc = entity_desc.strip('")')
        tmp_entity = entities_new.get(entity_name, [])
        if tmp_entity == []:
            entities_new[entity_name] = [[entity_type,entity_desc]]
        else:
            entities_new[entity_name].append([entity_type,entity_desc])

    # 描述合并
    res = []
    for entity_name,type_desc in tqdm(entities_new.items()):
        entity_type = type_desc[0][0]
        description_list = '\n'.join([desc for type,desc in type_desc])
        model_input = SUMMARIZE_PROMPT.replace('{entity_name}',entity_name)
        model_input = model_input.replace('{description_list}',description_list)

        entity_desc = chat(model_input,[],0.1)

        res.append([entity_name, entity_type, entity_desc])

    with open('../entities_new.pkl', 'wb') as f:
        pickle.dump(res,f)
